
#ifndef __DRIFT_H__
#define __DRIFT_H__

#include <iostream>
#include <fstream>
#include <vector>
#include <sstream>
#include <cstdlib>
#include <muParser/muParser.h>
#include <omd/treader.h>

// For 2D problems...

using std::string;
using std::vector;
using std::ifstream;
using std::istringstream;

namespace edrift {
  
#define FDCELL_EMPTY 0
#define FDCELL_OCCUPIED 1
#define FDCELL_PERIODIC 2
  
#define FD_MAX_VAR 50
  
  struct CellBoundary {
    char sym;
    // left right top bottom near far
    int  l,r,t,b,n,f;
  };
  
  enum eqntype {number, equ, compile}; 
    
  class DriftSolver {
    friend double read_table(double, double);
    friend double dread_table(double, double);
    friend double neig_cell_value(int, int, int);

    mu::Parser* mainequ;
    class DriftAlgorithm* algo;
    bool equ_ready; // set by equation_setup()
    bool link_updated; // unset by Link(), set by UpdateLink()

  public:

    // equation variables
    double var_x;
    double var_y;
    double var_z;
    double var_t;
    double var_lx;
    double var_ly;
    double var_lz;
    double var_lox;
    double var_loy;
    double var_loz;
    double var_hix; // changed in init constructor
    double var_hiy;
    double var_hiz;
    double var_dt;
    double var_dl;
    double var_cell;
    double var_step;
    double var_ix;
    double var_iy;
    double var_iz;
    double timecounter;
    int current_i;
    int current_j;
    int current_k;
    
    bool serial_process; // changed by check_expression(), used in AssignVar()

    double*  A;
    double*  T;
    double*  S;
    double*  Ext;
    double* dE;
    double* dS;
    double* dExt;
    double* sumE;
    double* sumS;
    double*  con;
    double*  cap;
    // limited to 5 couplers
    int ncoupler;
    class DriftCoupler* coupler[5];
    
    vector<double*> link;
    vector<string> linknm;
    
    int neqvar;
    struct {char name[32]; double value;} eqvar[FD_MAX_VAR];
    
    int me; // mpi rank
    int nproc; // number of proc
    int occup; // total # of occupied cells
    int ncp;
    int ncell;
    int *lim; // hi/lo limits of all procs
    int* counts;
    int* displ;
    int* sendnum;
    int* recvnum;
    int** nlist; // index to send
    int** glist; // index to receive
    double** ndata; // data to send
    double** gdata; // data to receive
    int** clist; // occupied cell list (not empty) of all procs
    
    vector<string> equname; // list of parser
    vector<string> equexpr; 
    vector<mu::Parser*> equpar;
    
    int dim,nx,ny,nz;
    int forced_periodic;
    int step;
    double timestep;
    double spacestep;
    double h2inv; // = 1/spacestep^2
    bool strict_courant;
    
    // map contains index to the coeficients
    int* map;
    
    int nbound;
    CellBoundary* bound;
    
    int search_symbol(char c);
    char map_symbol(int num);
    
    void get_cellcoord(int idx, int& i, int& j, int& k){
      int xy=nx*ny;
      int a=idx%xy;
      k=idx/xy;
      j=a/nx;
      i=a%nx;
    }
    
    void initiate_mem();
    void initiate_comm();
    void sync_data();
    void get_neigcell_index(int,int&,int&,int&,int&,int&,int&);
    void send_receive_list();
    void select_ncell(int,int,char**);
    bool mycell(int,int,int);
    double evaluate(string&);
    double evaluate(int,int,int);
    void check_expression(string&,int&,int&,eqntype&);
    
  public:
    int index(int i, int j, int k) {
      int idx=((k*nx*ny)+(i+j*nx));
      if(idx<0||idx>=dim) return -1;
      return idx;
    }

    void equation_setup();

    vector<omd::TableReader*> table;
    vector<string> tabname;
    
    string name;
    string str_cell; // fix value of cells
    string str_ext;  // coupling array
    string str_src;  // source term
    string str_cap;  // capacity
    string str_con;  // conductivity
    string str_var;  // variables
    string str_equ;  // compiled equations
    string str_link;  // external link pointer name
    
    bool has_cell;
    bool has_src;
    bool has_ext;
    bool has_con;
    bool has_cap;
    
    DriftSolver(int, int, int, string, string);
    virtual ~DriftSolver();
    
    void communicate(double* ptr);
    void collect(double* ptr);
    void collect(); // collect all
    void AssignArray(double*,int,int,int,double);
    void AssignArray(double*,string&);
    
    double* GetSourceVector() {return S;}
    double* GetExtVector() {return Ext;}
    double* GetPostSourceVector() {return dS;}
    double* GetPostExtVector() {return dExt;}
    double* GetVector(){return A;}
    double* GetConductivity(){return con;}
    double* GetCapacity(){return cap;}
    double* GetDelta(){return dE;}
    double* GetSourceTally(){return sumS;}
    double* GetExtTally(){return sumE;}
    double GetTimestep(){return timestep;}
    double GetSpacestep(){return spacestep;}
    int  GetStep(){return step;}
    double GetTime(){return timecounter;}
    int GetDim(){return dim;}
    void GetDim(int& _nx, int& _ny, int& _nz){_nx=nx;_ny=ny;_nz=nz;}
    char GetMapSymbol(int idx){return map_symbol(map[idx]);}

    double GetSum(double*);
    double GetAverage(double*);
    
    void LinkCoupler(class DriftCoupler*);
    void UnlinkCoupler(){ncoupler=0;} // this unlink all coupler! Relink any needed coupler.
    
    void CopyMap(char*); // convert map array to character map
    
    void AddTable(string);
    void AssignTableVar();
    
    void SetSymbolMap(string str);

    void SetMap(const char*);
    void Remap(const char*,const char*);   // change completely cell map
    
    void EnableDelta();
    void EnableTally();
    void KeepSource();
    
    void AddEquVar(string, double);
    void SetEquVar(string str) {str_var=str;}
    void SetCompileEqu(string str){str_equ=str;}

    virtual void AssignEquVar();

    void CheckStrings(); // to check these folowing strings. must be called after assigning.

    void SetCell(string str){str_cell=str;has_cell=true;}
    void SetExt(string str){str_ext=str;has_ext=true;}
    void SetSource(string str){str_src=str;has_src=true;}
    void SetConductivity(string str){str_con=str;has_con=true;}
    void SetCapacity(string str){str_cap=str;has_cap=true;}
    
    void ClearCell(){str_cell="";has_cell=false;}
    void ClearExt(){str_ext="";has_ext=false;}
    void ClearSource(){str_src="";has_src=false;}
    void ClearConductivity(){str_con="";has_con=false;}
    void ClearCapacity(){str_cap="";has_cap=false;}
    
    virtual void AssignCell(){if(str_cell.size())AssignArray(A,str_cell);}
    virtual void AssignExt(){if(str_ext.size())AssignArray(Ext,str_ext);}
    virtual void AssignSource(){if(str_src.size())AssignArray(S,str_src);}
    virtual void AssignConductivity(){if(str_con.size())AssignArray(con,str_con);}
    virtual void AssignCapacity(){if(str_cap.size())AssignArray(cap,str_cap);}
    
    void SetTimestep(double dt);
    void SetSpacestep(double h);
    
    void Inspect();
    void Step(int n);    
    void Dump(std::ostream& ofl);
    void Dump(const string fname);
    void ForcePeriodic(int bitmask) {forced_periodic=bitmask;} // (00000ZYX)b
    
    int GetNCP(){return ncp;} // # of local cell
    int* GetCellList(){return clist[me];}
    
    void SetLink(string str) {str_link=str;}
    void UpdateLink();
    void Link(string, double*);

    int coordtocell(double x, double y, double z);
    
    void DensityMap(
                    double* dens,
                    const char* syms,
                    double* limits,
                    const char* notouch=NULL);
    
    void MapToArray(double* dmap);
    
    void update_lox();
    void update_loy();
    void update_loz();
    void update_hix();
    void update_hiy();
    void update_hiz();
    
    std::ofstream wfile;
    void warning(string);
    void warning_file(string);
    
    // FIXME! to check all cell and give new timestep
    // factor: the factor of the criterion
    // change_dt: only calculate of put new value on timestep
    double CourantAdapt(double factor=1.0, bool change_dt=false);
    
    void SetAlgorithm(DriftAlgorithm* al);
    void ClearAlgorithm();
  
  };
  
  class DriftAlgorithm {
  public:
    DriftSolver* solver;
    DriftAlgorithm(){solver=NULL;}
    virtual void init(DriftSolver* slv){solver=slv;};
    virtual void exec();
  };
  
  // to be inherited by any coupling algorithms

  class DriftCoupler {
  public:
    DriftSolver* solver;
    DriftCoupler(){solver=NULL;}
    virtual void init(DriftSolver* slv){solver=slv;}
    virtual void assign(){};
    virtual void head(){};
    virtual void tail(){};
  };
  
}

#endif
